#!/usr/bin/env python

import PDZ_predictor as predict
import read_validation_files
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import sqlite3
import random
import csv
import picksdr
import sys

def ne_interaction():

    #create/check fasta file for each PDZ domain
    read_validation_files.read_domain_file('validate_data/pdzdomains.csv')
    
    #load peptide sequence
    pep = read_validation_files.read_peptide_file('validate_data/validation_peptides.csv')
    
    inter = read_validation_files.read_interaction_file('validate_data/interactions.csv', pep)
    output = list()
    for i in inter:
        if inter[i] > 0:
            #positive binding
            pass
        else:
            #negative binding
            output.append(i)
            
        #print i, inter[i], bind
    out2 = dict()
    for (i, j) in output:
        if out2.has_key(i):
            out2[i].append(j)
        else:
            out2[i] = [j]
            
    return out2

def score_negative(ne_binding = ne_interaction(), other=[]):
    
    #domain is misleading here, should really be protein
    scorelist = []
    fail_domain = []
    print "Total domain:", len(ne_binding)
    for domain in ne_binding:
        peptides = ne_binding[domain]
        try:
            seq_record = SeqIO.parse("validate_data/domains/"+domain+".fasta", "fasta").next()
        except IOError, e:
            print >> sys.stderr, e
        
        #print test
        result = predict.scoring (seq_record, peptides, "pdzvalid.db")
        if result is not False:
            scores, pdz_count = result
        else:
            fail_domain.append(domain)
            print >> sys.stderr, "Prediction can not be made for", domain
            
        for s in scores:
            if scores[s] is not False:
                for ss in scores[s]:
                    
                    scorelist.append(ss)
            else:
                fail_domain.append(domain)
                print >> sys.stderr, "Prediction can not be made for", domain, s
    #print fail_domain
    print "Total domain:", len(ne_binding)
    print "Total Failed:", len(fail_domain)
    return scorelist

def grouping(x):
    con = sqlite3.connect("pdzvalid.db")
    cur = con.cursor()
    subquery = """\
            SELECT DISTINCT domainID, sequence 
            FROM pdz
            """
    cur.execute(subquery)
    domain_ids = list()
    for row in cur:
        domain_ids.append(row)
    
    cur.close()
    con.close()
    
    #Divide them in to x groups.
    random.shuffle(domain_ids)
    vgroups = dict()
    
    for i in range(x):
        vgroups[i] = []
    i = 0
    while len(domain_ids) > 0:
        if i >= x:
            i = 0
        vgroups[i].append(domain_ids.pop())
        i += 1
    
    return vgroups

def data2file(x, i, scores, isbind, others = []):
    #x = x fold validation
    #i = the number ith group
    #isbind = 0 nonbinding, 1 binding
    if isbind == 0:
        print "Output to nonbind.dat."
        o = '-'.join(['{0}'.format(r) for r in others])
        outfile = 'vresult/{0}-nonbind{1}.dat'.format(x,o)
    else:
        o = '-'.join(['{0}'.format(r) for r in others])
        outfile = 'vresult/{0}-{1}-datafile{2}.dat'.format(x,i,o)
    with open (outfile, 'w') as fout:
        for k in scores:
            writer = csv.writer(fout, delimiter = ' ', lineterminator = '\n')
            writer.writerow([k, isbind])

def group2file(vgroups):
    x = len(vgroups)
    for group in vgroups:
        outfile = 'vresult/{0}vgroup{1}.txt'.format(x, group)
        
        with open (outfile, 'w') as fout:
            for k in vgroups[group]:
                writer = csv.writer(fout, delimiter = ' ', lineterminator = '\n')
                writer.writerow(k)

def loadgroup(x):
    vgroups = dict()
    for i in range(x):
        infile = 'vresult/{0}vgroup{1}.txt'.format(x, i)
        with open (infile, 'rb') as f:
            f = csv.reader(f, delimiter = ' ', lineterminator = '\n')
            for row in f:
                if vgroups.has_key(i):
                    vgroups[i].append((row[0],row[1]))
                else:
                    vgroups[i] = [(row[0], row[1])]
    
    return vgroups
    
                    
def make_groups(x):
    vgroups = grouping(x)
    group2file(vgroups)


def testgroup(fname):
    #just for testing
    domains = []
    with open(fname, 'rb') as f:
        d = csv.reader(f, delimiter = '\t', lineterminator = '\n')
        for i in d:
            domains.append(i)
    return domains



def testvalidation():
    domains = testgroup('validate_data/tgroup1.txt')

    #wh = raw_input()
    scorelist = []
    masklist = []
    for p in domains:
        masklist.append(p[0])
    for did in domains:
        print did
        con = sqlite3.connect("pdzvalid.db")
        cur = con.cursor()
        subquery = """\
                SELECT DISTINCT seq 
                FROM peptide_pdz
                where domainID = '{0}'
                """.format(did[0])
        cur.execute(subquery)
        peptides = list()
        for (row,) in cur:
            peptides.append(row)
        
        cur.close()
        con.close()
        #print peptides
        pseq = did[1]
        seq = predict.seq2seqrecord(pseq)
        #masklist = []
        #print did
        print seq, peptides
        result = predict.scoring (seq, peptides, "pdzvalid.db", masklist)        
        if result is not False:
            scores, pdz_count = result

        else:
            print "Prediction can not be made."
            scores = {-1:-1}
        for s in scores:
            for ss in scores[s]:
                scorelist.append(ss)
    print scorelist


    
    
def xvalidation(x, vgroups = False, others = []):
    
    if vgroups is False:
        vgroups = grouping(x)
    failed_domains = set()
    suceed_domains = set()
    for i in range(x):
        #i is the test set, others are training sets
        #get peptides for domains in i
        scorelist = []
        masklist = []
        
        for p in vgroups[i]:
            masklist.append(p[0])
        for did in vgroups[i]:
            #print did
            con = sqlite3.connect("pdzvalid.db")
            cur = con.cursor()
            subquery = """\
                    SELECT DISTINCT seq 
                    FROM peptide_pdz
                    where domainID = '{0}'
                    """.format(did[0])
            cur.execute(subquery)
            peptides = list()
            for (row,) in cur:
                peptides.append(row)
            
            cur.close()
            con.close()
            #print peptides
            pseq = did[1]
            seq = predict.seq2seqrecord(pseq)
            #masklist = []
            #print did
            result = predict.scoring (seq, peptides, "pdzvalid.db", masklist)        
            if result is not False:
                scores, pdz_count = result

            else:
                print >> sys.stderr,  "Prediction can not be made for ", did[0]
                scores = {-1:False}
            for s in scores:
                if scores[s] is not False:
                    for ss in scores[s]:
                        scorelist.append(ss)
                        suceed_domains.add(did[0])
                else:
                    print >> sys.stderr, "Prediction can not be made for ", did[0]
                    failed_domains.add(did[0])
                    #scores = {-1:-1}    
        #print sum(scorelist)/len(scorelist)
        data2file(x, i, scorelist, 1, others)
        print "Failed Domains:", len(failed_domains)
        print "TOTAL Domains:", len(suceed_domains) + len(failed_domains)

def cross_validation(x):
    tmp_sdr = []
    submrange = ['blosum62', 'blosum80'] #substitution matrix
    distrange = [4,5,6,8, 10] #contact distance
    cptrange = [0.2, 0.4,0.6,0.8] #min contact res % in SDR column
    sptrange = [0.2, 0.4,0.6,0.8] #max similar res % in SDR column
    dbname = "pdzvalid.db"
    contact_percent = 0.6
    similar_percent = 0.8
    subm = 'blosum62'
    dist = 8
    for contact_percent in cptrange:
        for similar_percent in sptrange:
            for subm in submrange:
                for dist in distrange:
                    
    
                    print "#-------------------------------------------------#"
                    print "#Testing parameters#"
                    print "Substitution Matrix:",subm
                    print "Max. Similar Res%:", similar_percent
                    print "Min. Contact Res%:", contact_percent
                    print "Max. Contact Distance:", dist
                
                    #conditions:
                    #contact percentage, similarity percentager, sub matrix, distance
                    vgroups = loadgroup(x)
                    flag_sdr = picksdr.select_sdr3(contact_percent, similar_percent,
                                subm, dist, dbname)
                    
                    #print flag_sdr
                    if flag_sdr in tmp_sdr:
                        pass
                    else:
                        tmp_sdr.append(flag_sdr)
                        if flag_sdr is not False:
                            picksdr.dbinsert(flag_sdr, dbname)
                            scorelist = score_negative(ne_interaction(), [contact_percent, similar_percent, subm, dist])
                            data2file(x, 0, scorelist, 0, [contact_percent, similar_percent, subm, dist])
                            xvalidation(x, vgroups, [contact_percent, similar_percent, subm, dist])
                        else:
                            print "ERROR: SDR Fail."

def sdr_test():
    #generate sdrs with different parameters
    #ignore the duplications, leave only the unique sdr sets (represent by one parameters)
    tmp_sdr = []
    tmp_par = []
    submrange = ['blosum62', 'blosum80'] #substitution matrix
    distrange = [4,5,6,8, 10] #contact distance
    cptrange = [0.2, 0.4,0.6,0.8] #min contact res % in SDR column
    sptrange = [0.2, 0.4,0.6,0.8] #max similar res % in SDR column
    dbname = "pdzvalid.db"
    for contact_percent in cptrange:
        for similar_percent in sptrange:
            for subm in submrange:
                for dist in distrange:

                    flag_sdr = picksdr.select_sdr3(contact_percent, similar_percent,
                                subm, dist, dbname)
                    
                    #print flag_sdr
                    if flag_sdr in tmp_sdr:
                        pass
                    else:
                        tmp_par.append((contact_percent, similar_percent, subm, dist, dbname))
                        tmp_sdr.append(flag_sdr)
    return tmp_sdr, tmp_par
    
if __name__ == '__main__':
    #AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    '''
    q_seq = "TVEIIRELSDALGISIAGGKGSPLGDIPIFIAMIQANGVAARTQKLKVGDRIVSINGQPLDGLSHTDAVNLLKNAFGRIILQVVADT"
    #seq = fasta2seqrecord("test.fasta")
    seq = predict.seq2seqrecord(q_seq)
    peptides = ["TKVLV","KKTTV"]
    scores, pdz_count = predict.scoring (seq, peptides, "pdzvalid.db")
    print pdz_count, scores
    '''
    
    #testvalidation()
    #flag_sdr = picksdr.select_sdr3(0.2, 0.8, 'blosum62', 5, 'pdzvalid.db')
                    
    #print flag_sdr
    #print flag_sdr
    
    #if flag_sdr is not False:
    #    picksdr.dbinsert(flag_sdr, 'pdzvalid.db')
    
        
    sdr_sets, par = sdr_test()
    for s in par:
        print '-------------'
        print s
    #make_groups(8)
    #cross_validation(176)